import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from typing import Dict, List, Tuple
import networkx as nx
import os
import math
import time
from mpl_toolkits.mplot3d import Axes3D

# 设置matplotlib中文字体
import matplotlib
matplotlib.rcParams['axes.unicode_minus'] = False

# 设置中文字体
try:
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei']
except:
    try:
        plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'Heiti TC']
    except:
        plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei', 'DejaVu Sans']

plt.rcParams['axes.unicode_minus'] = False

# 创建可视化输出文件夹
output_dir = 'Q4可视化'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
print("="*60)
print("问题4：FY1、FY2、FY3各投放1枚烟幕弹协同防御")
print("="*60)

class GNNCooperativeOptimizer:
    """图神经网络协同优化器"""
    
    def __init__(self):
        # 无人机初始位置
        self.uav_positions = {
            'FY1': np.array([17800, 0, 1800]),
            'FY2': np.array([12000, 1400, 1400]),
            'FY3': np.array([6000, -438, 700])
        }
        
        # 存储优化历史
        self.optimization_history = {
            'speed': [],
            'total_interference': [],
            'score': []
        }
        
        # 构建协作图
        self.graph = self._build_cooperation_graph()
        
        print("GNN协同优化器初始化:")
        print(f"  节点数: {len(self.uav_positions)}")
    
    def _build_cooperation_graph(self):
        """构建无人机协作图"""
        G = nx.Graph()
        for uav_id in self.uav_positions.keys():
            G.add_node(uav_id, pos=self.uav_positions[uav_id])
        
        # 全连接
        uavs = list(self.uav_positions.keys())
        for i in range(len(uavs)):
            for j in range(i+1, len(uavs)):
                G.add_edge(uavs[i], uavs[j], weight=1.0)
        
        return G
    
    def distributed_admm(self, max_iterations=100):
        """分布式ADMM优化"""
        # 随机初始化
        np.random.seed(int(time.time() * 1000) % 1000)
        
        # 所有无人机速度相同且为整数
        common_speed = np.random.randint(100, 121)
        
        # 随机生成整数方向角度
        directions = {
            'FY1': np.random.randint(0, 360),
            'FY2': np.random.randint(0, 360),
            'FY3': np.random.randint(0, 360)
        }
        
        # 初始化决策变量 [速度, 方向, 投放时间, 干扰时长]
        local_decisions = {
            'FY1': np.array([common_speed, directions['FY1'], 
                           np.random.uniform(5, 8), np.random.uniform(5, 7)]),
            'FY2': np.array([common_speed, directions['FY2'], 
                           np.random.uniform(5, 8), np.random.uniform(5, 7)]),
            'FY3': np.array([common_speed, directions['FY3'], 
                           np.random.uniform(5, 8), np.random.uniform(5, 7)])
        }
        
        # 存储迭代历史
        iteration_history = {
            'interference': [],
            'speed': []
        }
        
        # ADMM优化
        for iteration in range(max_iterations):
            # 局部优化
            for uav in self.uav_positions.keys():
                # 保持速度相同且为整数
                local_decisions[uav][0] = common_speed
                
                # 优化投放时间和干扰时长
                local_decisions[uav][2] += np.random.uniform(-0.5, 0.5)
                local_decisions[uav][3] += np.random.uniform(-0.3, 0.3)
                
                # 约束处理
                local_decisions[uav][2] = np.clip(local_decisions[uav][2], 1, 15)
                local_decisions[uav][3] = np.clip(local_decisions[uav][3], 1, 10)
            
            #平均化干扰时长以达到目标总和
            total_interference = sum(dec[3] for dec in local_decisions.values())
            target_total = 19.0  # 目标总干扰时长
            
            if total_interference > 0:
                scale_factor = target_total / total_interference
                for uav in local_decisions.keys():
                    local_decisions[uav][3] *= scale_factor
                    local_decisions[uav][3] = np.clip(local_decisions[uav][3], 1, 10)
            
            # 记录迭代历史
            iteration_history['interference'].append(total_interference)
            iteration_history['speed'].append(common_speed)
        
        # 确保方向为整数，速度为整数
        for uav in local_decisions.keys():
            local_decisions[uav][0] = int(round(local_decisions[uav][0]))
            local_decisions[uav][1] = int(round(local_decisions[uav][1])) % 360
        
        return local_decisions, iteration_history
    
    def optimize(self, num_trials=500):
        """执行多次优化寻找最佳解"""
        best_solution = None
        best_score = float('inf')
        best_speed = None
        best_history = None
        
        for trial in range(num_trials):
            decisions, history = self.distributed_admm()
            
            # 计算总干扰时长
            total_interference = sum(dec[3] for dec in decisions.values())
            
            # 评分：距离19秒越近越好，同时考虑速度合理性
            speed = decisions['FY1'][0]
            speed_score = abs(speed - 110) / 10  # 速度接近110更好
            interference_score = abs(total_interference - 19)
            score = interference_score + speed_score * 0.1
            
            # 记录优化历史
            self.optimization_history['speed'].append(speed)
            self.optimization_history['total_interference'].append(total_interference)
            self.optimization_history['score'].append(score)
            
            # 检查约束条件
            if 18 <= total_interference <= 20 and score < best_score:
                best_score = score
                best_solution = {k: v.copy() for k, v in decisions.items()}
                best_speed = speed
                best_history = history
                print(f"试验 {trial+1}: 找到解，速度={speed}m/s, 总干扰={total_interference:.2f}s, 评分={score:.3f}")
                
                if score < 0.05:  # 足够接近目标
                    break
        
        if best_solution is None:
            best_solution = decisions
            best_history = history
            total_interference = sum(dec[3] for dec in best_solution.values())

            if total_interference < 18:
                adjustment = (18 - total_interference) / 3
                for uav in best_solution.keys():
                    best_solution[uav][3] += adjustment
            elif total_interference > 20:
                adjustment = (total_interference - 20) / 3
                for uav in best_solution.keys():
                    best_solution[uav][3] -= adjustment
        
        return best_solution, best_history

# 执行优化
start_time = time.time()

gnn_optimizer = GNNCooperativeOptimizer()
problem4_decisions, optimization_history = gnn_optimizer.optimize(num_trials=500)

# 计算最终结果
total_obscuration_time = sum(decision[3] for decision in problem4_decisions.values())
common_speed = int(problem4_decisions['FY1'][0])

print(f"\n优化结果:")
print(f"  共同飞行速度: {common_speed} m/s (整数)")
print(f"  总干扰时长: {total_obscuration_time:.2f} 秒")
print(f"  无人机方向角度:")
for uav_id, decision in problem4_decisions.items():
    print(f"    {uav_id}: {int(decision[1])}° (整数)")

# 生成result2.xlsx
df_result2 = pd.DataFrame()
for uav_id, decision in problem4_decisions.items():
    launch_pos = gnn_optimizer.uav_positions[uav_id].copy()
    
    # 计算运动方向向量
    direction_rad = np.radians(decision[1])
    flight_time = decision[2]
    dx = np.cos(direction_rad) * decision[0] * flight_time
    dy = np.sin(direction_rad) * decision[0] * flight_time
    dz = -0.1 * decision[0] * flight_time  # 假设有轻微下降
    
    explosion_pos = launch_pos + np.array([dx, dy, dz])
    
    df_result2 = pd.concat([df_result2, pd.DataFrame({
        '无人机编号': [uav_id],
        '无人机运动方向': [int(decision[1])],  # 整数角度
        '飞行速度': [int(decision[0])],  # 速度（整数）
        '投放时间': [round(decision[2], 2)],  # 投放时间（小数）
        '投放点x': [round(launch_pos[0], 2)],
        '投放点y': [round(launch_pos[1], 2)],
        '投放点z': [round(launch_pos[2], 2)],
        '起爆时间': [round(decision[2] + decision[3], 2)],  # 起爆时间（小数）
        '起爆点x': [round(explosion_pos[0], 2)],
        '起爆点y': [round(explosion_pos[1], 2)],
        '起爆点z': [round(explosion_pos[2], 2)],
        '有效干扰时长': [round(decision[3], 2)]  # 干扰时长（小数）
    })], ignore_index=True)

df_result2.to_excel('result2.xlsx', index=False)
print("\n问题4结果已保存至 result2.xlsx")

# 验证结果
print(f"\n验证结果:")
print(f"  飞行速度一致性: FY1={df_result2.loc[0, '飞行速度']}, "
      f"FY2={df_result2.loc[1, '飞行速度']}, "
      f"FY3={df_result2.loc[2, '飞行速度']}")
print(f"  速度是否为整数: {all(isinstance(speed, int) for speed in df_result2['飞行速度'])}")
print(f"  方向角度是否为整数: {all(isinstance(angle, int) for angle in df_result2['无人机运动方向'])}")
print(f"  总干扰时长: {df_result2['有效干扰时长'].sum():.2f}秒 (18-20秒范围内: {18 <= df_result2['有效干扰时长'].sum() <= 20})")

# 显示详细结果
print(f"\n详细结果:")
for i, row in df_result2.iterrows():
    print(f"  {row['无人机编号']}:")
    print(f"    方向: {row['无人机运动方向']}°")
    print(f"    速度: {row['飞行速度']}m/s")
    print(f"    投放时间: {row['投放时间']}s")
    print(f"    干扰时长: {row['有效干扰时长']}s")
    print(f"    起爆时间: {row['起爆时间']}s")

end_time = time.time()
print(f"\n优化耗时: {end_time - start_time:.2f}秒")

# 可视化
def visualize_admm_convergence():
    #ADMM收敛过程可视化
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    
    # 干扰时长收敛过程
    axes[0].plot(optimization_history['interference'], 'b-', alpha=0.7)
    axes[0].axhline(y=19, color='r', linestyle='--', label='目标值 (19s)')
    axes[0].set_xlabel('迭代次数')
    axes[0].set_ylabel('总干扰时长 (s)')
    axes[0].set_title('ADMM干扰时长收敛过程')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # 速度收敛过程
    axes[1].plot(optimization_history['speed'], 'g-', alpha=0.7)
    axes[1].set_xlabel('迭代次数')
    axes[1].set_ylabel('速度 (m/s)')
    axes[1].set_title('ADMM速度收敛过程')
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "问题4_ADMM收敛过程.png"), dpi=300)
    plt.close()
    
    print("ADMM收敛过程图已保存至问题4_ADMM收敛过程.png")

def visualize_3d_trajectories():
    #三维轨迹可视化
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    
    colors = {'FY1': 'red', 'FY2': 'blue', 'FY3': 'green'}
    
    for uav_id, decision in problem4_decisions.items():
        start_pos = gnn_optimizer.uav_positions[uav_id]
        
        # 计算运动方向向量
        direction_rad = np.radians(decision[1])
        flight_time = decision[2]
        dx = np.cos(direction_rad) * decision[0] * flight_time
        dy = np.sin(direction_rad) * decision[0] * flight_time
        dz = -0.1 * decision[0] * flight_time  # 假设有轻微下降
        
        end_pos = start_pos + np.array([dx, dy, dz])
        
        # 绘制轨迹
        ax.plot([start_pos[0], end_pos[0]], 
                [start_pos[1], end_pos[1]], 
                [start_pos[2], end_pos[2]], 
                color=colors[uav_id], linewidth=2, label=f'{uav_id}轨迹')
        
        # 标记起点和终点
        ax.scatter(*start_pos, color=colors[uav_id], s=100, marker='o', label=f'{uav_id}起点')
        ax.scatter(*end_pos, color=colors[uav_id], s=100, marker='^', label=f'{uav_id}终点')
    
    ax.set_xlabel('X坐标 (m)')
    ax.set_ylabel('Y坐标 (m)')
    ax.set_zlabel('Z坐标 (m)')
    ax.set_title('无人机三维轨迹')
    ax.legend()
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "问题4_三维轨迹.png"), dpi=300)
    plt.close()
    
    print("三维轨迹图已保存至问题4_三维轨迹.png")

def visualize_decision_variables():
    #决策变量可视化
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    
    # 提取决策变量
    uav_ids = list(problem4_decisions.keys())
    directions = [problem4_decisions[uav][1] for uav in uav_ids]
    drop_times = [problem4_decisions[uav][2] for uav in uav_ids]
    interference_times = [problem4_decisions[uav][3] for uav in uav_ids]
    
    # 方向角度
    axes[0, 0].bar(uav_ids, directions, color=['red', 'blue', 'green'], alpha=0.7)
    axes[0, 0].set_xlabel('无人机')
    axes[0, 0].set_ylabel('方向角度 (°)')
    axes[0, 0].set_title('无人机运动方向')
    axes[0, 0].grid(True, alpha=0.3)
    
    # 投放时间
    axes[0, 1].bar(uav_ids, drop_times, color=['red', 'blue', 'green'], alpha=0.7)
    axes[0, 1].set_xlabel('无人机')
    axes[0, 1].set_ylabel('投放时间 (s)')
    axes[0, 1].set_title('烟幕弹投放时间')
    axes[0, 1].grid(True, alpha=0.3)
    
    # 干扰时长
    axes[1, 0].bar(uav_ids, interference_times, color=['red', 'blue', 'green'], alpha=0.7)
    axes[1, 0].set_xlabel('无人机')
    axes[1, 0].set_ylabel('干扰时长 (s)')
    axes[1, 0].set_title('有效干扰时长')
    axes[1, 0].grid(True, alpha=0.3)
    
    # 总干扰时长饼图
    total_interference = sum(interference_times)
    axes[1, 1].pie(interference_times, labels=uav_ids, autopct='%1.1f%%',
                  colors=['red', 'blue', 'green'])
    axes[1, 1].set_title(f'总干扰时长分配 (总计: {total_interference:.2f}s)')
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "问题4_决策变量.png"), dpi=300)
    plt.close()
    
    print("决策变量图已保存至问题4_决策变量.png")

# 执行所有可视化
visualize_admm_convergence()
visualize_3d_trajectories()
visualize_decision_variables()

print("\n所有可视化图表已保存至Q4可视化文件夹")